import sys
import os

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import torch
import yaml
import cv2
from modules.xfeat import XFeat
from gl3d_pose_est_mnn import GL3DPoseMNNBenchmark
from dual_softmax_matcher import DualSoftMaxMatcher


def get_best_device(verbose=False):
    device = torch.device('cpu')
    if torch.cuda.is_available():
        device = torch.device('cuda')
    elif torch.backends.mps.is_available():
        device = torch.device('mps')
    else:
        device = torch.device('cpu')
    if verbose: print(f"Fastest device found is: {device}")
    return device


if __name__ == "__main__":
    matcher = DualSoftMaxMatcher()
    weight_path = "/media/liyuke/share/AAA/part2/xfeat/weights/xfeat.pt"
    model = XFeat(weights=weight_path)
    # init dataset
    gl3d = GL3DPoseMNNBenchmark(data_root="/media/liyuke/share/GL3D/")
    gl3d.benchmark(desc_model=model,
                   matcher_model=matcher)
